import os
import core

# 禁用gpu
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"


def zhuTi_Predict(data):
    """
    主题预测
    :param data:待预测数据
    :return:7类主题（数组）
    """
    tokenizer_addr = '.\\data\\tokenizer_data\\zhuTi_Tokenizer.pickle'
    model_addr = '.\\data\\model_data\\zhu_ti.h5'

    res = core.predict(data, tokenizer_addr, model_addr)
    out_data = {
        0: '空间',
        1: '动力',
        2: '油耗',
        3: '操控',
        4: '内饰',
        5: '舒适性',
        6: '外观'
    }
    out = []
    for i in res:
        out.append(out_data.get(i))
    return out


def qingGan_Predict(data):
    """
    情感预测
    :param data:待预测数据
    :return:2类情感（数组）
    """
    tokenizer_addr = '.\\data\\tokenizer_data\\qingGan_Tokenizer.pickle'
    model_addr = '.\\data\\model_data\\qing_gan.h5'
    res = core.predict(data, tokenizer_addr, model_addr)
    out_data = {
        0: '不满意',
        1: '满意'
    }
    out = []
    for i in res:
        out.append(out_data.get(i))
    return out


def predict(data):
    """
    对数据预测
    :param data: 待预测数据
    :return: 预测结果（情感，主题 键值对数组）
    """
    # 提取文本
    train_data = []
    for each in data:
        train_data.append(each[1])

    qing_gan = qingGan_Predict(train_data)
    zhu_ti = zhuTi_Predict(train_data)

    i = 0
    result = []
    # 以集合数组传出
    while i < len(data):
        result.append({
            'Pid': data[i][0],
            'Emotional': qing_gan[i],
            'Theme': zhu_ti[i]
        })
        i += 1
    return result


if __name__ == '__main__':
    data = [[1, '行驶品质超赞，车内噪音更是特别的小。另外，超低的油耗，让其他的同事都很羡慕。'], [2, '行驶品质超赞，车内噪音更是特别的小。另外，超低的油耗，让其他的同事都很羡慕。']]
    print(predict(data))
